/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: alltoall pipeline 抽象基类实现
 * Author: zhujiale
 * Create: 2024-03-15
 */

#include "allltoall_pipeline_base.h"

namespace hccl {
AlltoallPipelineBase::AlltoallPipelineBase(
    const HcclDispatcher dispatcher,
    const std::vector<SendRecvInfo> &allMeshAggregationSendRecvInfo,
    HcclWorkflowMode workMode): AlgTemplateBase(dispatcher),
    allMeshAggregationSendRecvInfo_(allMeshAggregationSendRecvInfo), workMode_(workMode)
{}

AlltoallPipelineBase::~AlltoallPipelineBase() {}

HcclResult AlltoallPipelineBase::Prepare(u32 userRank, A2aPipelineMemory A2aPipelineMemory,
    std::unique_ptr<CommBase> &commLevel0, std::unique_ptr<CommBase> &commLevel1,
    Stream &mainStream, std::vector<Stream> &subStream,
    std::vector<std::shared_ptr<LocalNotify>> &notifyMain, std::vector<std::shared_ptr<LocalNotify>> &notifySub)
{
    localSendRecvInfo_ = allMeshAggregationSendRecvInfo_[userRank];

    inputMem_ = A2aPipelineMemory.userInput;
    outputMem_ = A2aPipelineMemory.userOutput;
    scratchMem_ = A2aPipelineMemory.scratchMem;
    cclIn_ = A2aPipelineMemory.cclInBuffer;
    cclOut_ = A2aPipelineMemory.cclOutBuffer;

    intraRankSize_ = commLevel0->RankSize();
    interRankSize_ = commLevel1->RankSize();
    groupRankSize_ = intraRankSize_ * interRankSize_;

    userRank_ = userRank;
    intraRankId_ = commLevel0->Rank();
    interRankId_ = commLevel1->Rank();

    meshRankStart_ = userRank - intraRankId_;
    meshRankEnd_ = meshRankStart_ + intraRankSize_ - 1;

    mainStream_ = mainStream;
    subStream_ = subStream;
    streamNotifyMain_ = notifyMain;
    streamNotifySub_ = notifySub;

    intraLinks_ = commLevel0->TransportInfo();
    interLinks_ = commLevel1->TransportInfo();

    HCCL_DEBUG("[AlltoallPipelineBase]streamNum[%u], streamNotifyMainNum[%u], streamNotifySubNum[%u]",
        subStream_.size(), streamNotifyMain_.size(), streamNotifySub_.size());
    HCCL_DEBUG("[AlltoallPipelineBase]interLinksNum[%u], intraLinksNum[%u]",
        interLinks_.size(), intraLinks_.size());

    return HCCL_SUCCESS;
}

HcclResult AlltoallPipelineBase::Prepare(u32 userRank, A2aPipelineMemory A2aPipelineMemory,
    const SubCommInfo &level0CommInfo, const SubCommInfo &level1CommInfo,
    Stream &mainStream, std::vector<Stream> &subStream,
    std::vector<std::shared_ptr<LocalNotify>> &notifyMain, std::vector<std::shared_ptr<LocalNotify>> &notifySub)
{
    localSendRecvInfo_ = allMeshAggregationSendRecvInfo_[userRank];

    inputMem_ = A2aPipelineMemory.userInput;
    outputMem_ = A2aPipelineMemory.userOutput;
    scratchMem_ = A2aPipelineMemory.scratchMem;
    cclIn_ = A2aPipelineMemory.cclInBuffer;
    cclOut_ = A2aPipelineMemory.cclOutBuffer;

    intraRankSize_ = level0CommInfo.localRankSize;
    interRankSize_ = level1CommInfo.localRankSize;
    groupRankSize_ = intraRankSize_ * interRankSize_;

    userRank_ = userRank;
    intraRankId_ = level0CommInfo.localRank;
    interRankId_ = level1CommInfo.localRank;

    meshRankStart_ = userRank - intraRankId_;
    meshRankEnd_ = meshRankStart_ + intraRankSize_ - 1;

    mainStream_ = mainStream;
    subStream_ = subStream;
    streamNotifyMain_ = notifyMain;
    streamNotifySub_ = notifySub;

    intraLinks_ = level0CommInfo.links;
    interLinks_ = level1CommInfo.links;

    HCCL_DEBUG("[AlltoallPipelineBase]streamNum[%u], streamNotifyMainNum[%u], streamNotifySubNum[%u]",
        subStream_.size(), streamNotifyMain_.size(), streamNotifySub_.size());
    HCCL_DEBUG("[AlltoallPipelineBase]interLinksNum[%u], intraLinksNum[%u]", interLinks_.size(), intraLinks_.size());

    return HCCL_SUCCESS;
}

HcclResult AlltoallPipelineBase::CheckResourceValid()
{
    if (workMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) {
        CHK_PRT_RET(cclIn_.size() != cclOut_.size(),
            HCCL_ERROR("[AlltoallPipelineBase][CheckResourceValid] cclIn mem and cclOut mem should be the same size, "
            "ScratchInputMem[%llu] ScratchOutputMem[%llu]", cclIn_.size(), cclOut_.size()),
            HCCL_E_MEMORY);
    }
    CHK_PRT_RET(subStream_.size() < intraRankSize_ || streamNotifyMain_.size() < intraRankSize_ ||
        streamNotifySub_.size() < intraRankSize_, HCCL_DEBUG("[AlltoallPipelineBase][CheckResourceValid] "
        "stream resource not enough, num sub stream[%llu], num notify main signal[%llu] num notify sub signal[%llu], "
        "should be more than or equal to intraRankSize %llu", subStream_.size(), streamNotifyMain_.size(),
        streamNotifySub_.size(), intraRankSize_), HCCL_E_UNAVAIL);
    return HCCL_SUCCESS;
}

// alltoall 系列算法抽象行为应该都可以分为第一次发送前的数据准备，中间的每一步同步发送，以及本地数据搬移
HcclResult AlltoallPipelineBase::RunAsync()
{
    CHK_RET(CheckResourceValid());
    CHK_RET(PreProcess());
    for (u32 step = 0, numStep = CalcInterNumSteps(); step < numStep; step++) {
        CHK_RET(PipelineSend(step, step == (numStep - 1)));
    }
    CHK_RET(PostProcess());
    return HCCL_SUCCESS;
}

std::string AlltoallPipelineBase::GetCurrClassName()
{
    std::string className = typeid(*this).name();
    if (className.find("class") != className.npos) {
        size_t classNamePrefixLen = 6;
        className = className.substr(classNamePrefixLen);
    }
    return className;
}

std::string AlltoallPipelineBase::GetStreamIndexString()
{
    std::string res = "";
    for (auto& info : intraStreamInfo_) {
        res += std::to_string(info.first) + ", ";
    }
    return res;
}

HcclResult AlltoallPipelineBase::NotifyInterStreamStart()
{
    CHK_RET(LocalNotify::Post(mainStream_, dispatcher_, streamNotifySub_[intraRankId_],
        INVALID_VALUE_STAGE));
    CHK_RET(LocalNotify::Wait(subStream_[intraRankId_], dispatcher_, streamNotifySub_[intraRankId_],
        INVALID_VALUE_STAGE));
    HCCL_DEBUG("[%s][NotifyInterStreamStart] userRank %u, interRank %u, "
        "intraRank %u, main stream notify sdma stream %s", GetCurrClassName().c_str(),
        userRank_, interRankId_, intraRankId_, GetStreamIndexString().c_str());
    return HCCL_SUCCESS;
}

HcclResult AlltoallPipelineBase::WaitInterStreamFinish()
{
    CHK_RET(LocalNotify::Post(subStream_[intraRankId_], dispatcher_, streamNotifyMain_[intraRankId_],
        INVALID_VALUE_STAGE));
    CHK_RET(LocalNotify::Wait(mainStream_, dispatcher_, streamNotifyMain_[intraRankId_],
        INVALID_VALUE_STAGE));
    HCCL_DEBUG("[%s][WaitInterStreamFinish] userRank %u, interRank %u, intraRank %u, "
        "main stream notify sdma stream %s", GetCurrClassName().c_str(), userRank_, interRankId_,
        intraRankId_, GetStreamIndexString().c_str());
    return HCCL_SUCCESS;
}

// 主流只需要通知当前子步骤需要收发数据的 SDMA 流，减少同步开销
HcclResult AlltoallPipelineBase::NotifyIntraStreamStart()
{
    for (auto& sdmaInfo : intraStreamInfo_) {
        u32 streamIndex = sdmaInfo.first;
        CHK_RET(LocalNotify::Post(mainStream_, dispatcher_, streamNotifySub_[streamIndex],
            INVALID_VALUE_STAGE));
        CHK_RET(LocalNotify::Wait(subStream_[streamIndex], dispatcher_, streamNotifySub_[streamIndex],
            INVALID_VALUE_STAGE));
    }
    HCCL_DEBUG("[%s][NotifyIntraStreamStart] userRank %u, interRank %u, "
        "intraRank %u, main stream notify sdma stream %s", GetCurrClassName().c_str(),
        userRank_, interRankId_, intraRankId_, GetStreamIndexString().c_str());
    return HCCL_SUCCESS;
}

HcclResult AlltoallPipelineBase::WaitIntraStreamFinish()
{
    for (auto& sdmaInfo : intraStreamInfo_) {
        u32 streamIndex = sdmaInfo.first;
        CHK_RET(LocalNotify::Wait(mainStream_, dispatcher_, streamNotifyMain_[streamIndex],
            INVALID_VALUE_STAGE));
        CHK_RET(LocalNotify::Post(subStream_[streamIndex], dispatcher_, streamNotifyMain_[streamIndex],
            INVALID_VALUE_STAGE));
    }
    HCCL_DEBUG("[%s][WaitIntraStreamFinish] userRank %u, interRank %u, "
        "intraRank %u, main stream wait sdma stream %s", GetCurrClassName().c_str(), userRank_,
        interRankId_, intraRankId_, GetStreamIndexString().c_str());
    return HCCL_SUCCESS;
}

} // namespace hccl